import json
import os

# 注意运行2次，train/val
# 修改1 json存储地址
# json_path = "D:\\Tools\\yolo\\dlb\\PCB_DATASET\\Annotations\\val.json"
json_path = "D:\\Tools\\yolo\\dlb\\PCB_DATASET\\Annotations\\train.json"
# # 修改2，保存位置yolo存储地址
# yolo_paths = "D:\\Tools\\yolo\\dlb\\PCB_DATASET\\labels\\val2017/"
yolo_paths = "D:\\Tools\\yolo\\dlb\\PCB_DATASET\\labels\\train2017/"

with open(json_path) as f:
    data = json.load(f)

imgs = {}
for img in data['images']:
    imgs[img['id']] = {
        'h': img['height'],
        'w': img['width'],
        'file_name': img['file_name'],
    }

tmp = ''
for anno in data['annotations']:
    print(imgs[anno['image_id']]['file_name'])

    if imgs[anno['image_id']] != tmp:
        txt_path = os.path.join(yolo_paths, imgs[anno['image_id']]['file_name'].split('.')[0] + '.txt')
        txt_file = open(txt_path, 'w')
        # xywh --> xywh(归一化)
        bbox = [anno['bbox'][0] / imgs[anno['image_id']]['w'],
                anno['bbox'][1] / imgs[anno['image_id']]['h'],
                anno['bbox'][2] / imgs[anno['image_id']]['w'],
                anno['bbox'][3] / imgs[anno['image_id']]['h']]

        cls_id = anno['category_id']
        # 保存
        txt_file.write(str(cls_id) + ' ' + " ".join([str(a) for a in bbox]) + "\n")  # 生成格式0 cx,cy,w,h
        tmp = imgs[anno['image_id']]

    else:
        # xywh --> xywh(归一化)
        bbox = [anno['bbox'][0] / imgs[anno['image_id']]['w'],
                anno['bbox'][1] / imgs[anno['image_id']]['h'],
                anno['bbox'][2] / imgs[anno['image_id']]['w'],
                anno['bbox'][3] / imgs[anno['image_id']]['h']]

        cls_id = anno['category_id']
        # 保存
        txt_file.write(str(cls_id) + ' ' + " ".join([str(a) for a in bbox]) + "\n")  # 生成格式0 cx,cy,w,h
